#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : cd
# @Email : 19688513@qq.com
# @Project : horizons-engine-pybroker
# @File : StockSignalsManager.py
# @Software: PyCharm
# @time: 2025/6/30 10:11
# @description : 股票信号管理器 - 处理信号生成、存储和验证
"""

import urllib3
import logging
import pandas as pd
import numpy as np
from tqdm import tqdm
from typing import Dict, List, Any, Callable, Optional
import sqlite3
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime

from SignalGenerator import SignalGenerator
from StockBasicInfoManager import StockBasicInfoManager
from TradeDateManager import TradeDateManager
from YearlyStockDataManager import YearlyStockDataManager

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class StockSignalsManager:
    def __init__(self, trade_calendar: TradeDateManager,
                 data_manager: YearlyStockDataManager,
                 signal_generator: SignalGenerator,
                 stock_manager: StockBasicInfoManager,
                 db_path='data/stock_signals.db'):
        """
        初始化股票信号处理类

        :param trade_calendar: 交易日历对象
        :param data_manager: 数据管理器对象
        :param signal_generator: 信号生成器对象
        :param db_path: 数据库文件路径
        """
        # 确保数据目录存在
        os.makedirs(os.path.dirname(db_path), exist_ok=True)

        self.trade_calendar = trade_calendar
        self.data_manager = data_manager
        self.signal_generator = signal_generator
        self.stock_manager = stock_manager
        self.db_path = db_path
        self.stock_info_map: Dict[str, Dict[str, str]] = {}
        self.batch_size = 100  # 批量处理大小
        self._init_db()

        # 检查并确保交易日历已加载
        if not trade_calendar._trade_dates:
            logging.warning("交易日历未加载数据，正在初始化...")
            trade_calendar.initialize_trade_dates()

        # 检查并确保股票信息已加载
        try:
            stock_count = len(stock_manager.get_stock_codes())
            if stock_count == 0:
                logging.warning("股票代码列表为空，尝试初始化数据库...")
                stock_manager.initialize_stock_database()
        except Exception as e:
            logging.error(f"初始化股票数据库失败: {e}")


    def _init_db(self):
        """初始化数据库表结构"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # 创建主信号表
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS stock_signals (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            stock_code TEXT NOT NULL,
            signal_date TEXT NOT NULL,
            market TEXT NOT NULL,
            signal_type TEXT,
            signal_data TEXT,
            decline_condition_met INTEGER DEFAULT 0,  -- 0:未满足, 1:满足
            buy_signal INTEGER DEFAULT 0,             -- 0:无信号, 1:买入信号
            sell_signal INTEGER DEFAULT 0,            -- 0:无信号, 1:卖出信号
            created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE(stock_code, signal_date, market)
        )
        ''')

        # 创建买入信号历史表
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS buy_signals_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            stock_code TEXT NOT NULL,
            buy_date TEXT NOT NULL,
            signal_data TEXT,
            created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')

        # 创建卖出信号历史表
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS sell_signals_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            stock_code TEXT NOT NULL,
            sell_date TEXT NOT NULL,
            buy_date TEXT NOT NULL,
            signal_data TEXT,
            created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')

        conn.commit()
        conn.close()

    def store_signals(self, signals: List[Dict], market: str):
        """
        存储信号到数据库
        """
        if not signals:
            return

        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        try:
            for signal in signals:
                # 准备数据
                stock_code = signal.get('stock_code')
                signal_date = signal.get('date')
                signal_type = signal.get('signal_type')
                signal_data = json.dumps(signal)  # 序列化为JSON字符串

                # 插入或更新记录
                cursor.execute('''
                INSERT OR REPLACE INTO stock_signals 
                (stock_code, signal_date, market, signal_type, signal_data)
                VALUES (?, ?, ?, ?, ?)
                ''', (stock_code, signal_date, market, signal_type, signal_data))

            conn.commit()
            logging.info(f"成功存储 {len(signals)} 条{market}市场信号到数据库")
        except Exception as e:
            logging.error(f"存储信号到数据库失败: {str(e)}")
        finally:
            conn.close()

    def get_signals_by_date(self, target_date: str, filter_condition: Optional[str] = None) -> Dict[str, List[Dict]]:
        """
        从数据库获取指定日期的信号，可过滤条件

        :param target_date: 目标日期
        :param filter_condition: 过滤条件 ('buy', 'sell', 'decline')
        :return: 包含深圳和上海市场信号的字典
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        sz_signals = []
        sh_signals = []

        # 构建查询条件
        condition_clause = ""
        if filter_condition:
            if filter_condition == 'buy':
                condition_clause = "AND buy_signal = 1"
            elif filter_condition == 'sell':
                condition_clause = "AND sell_signal = 1"
            elif filter_condition == 'decline':
                condition_clause = "AND decline_condition_met = 1"

        try:
            # 查询深圳市场信号
            cursor.execute(f'''
            SELECT signal_data FROM stock_signals 
            WHERE signal_date = ? AND market = ? {condition_clause}
            ''', (target_date, 'sz'))
            for row in cursor.fetchall():
                sz_signals.append(json.loads(row[0]))

            # 查询上海市场信号
            cursor.execute(f'''
            SELECT signal_data FROM stock_signals 
            WHERE signal_date = ? AND market = ? {condition_clause}
            ''', (target_date, 'sh'))
            for row in cursor.fetchall():
                sh_signals.append(json.loads(row[0]))

            logging.info(
                f"从数据库读取: 深圳市场{len(sz_signals)}条, 上海市场{len(sh_signals)}条 (过滤条件: {filter_condition})")
        except Exception as e:
            logging.error(f"从数据库获取信号失败: {str(e)}")
        finally:
            conn.close()

        return {
            'sz_signals': sz_signals,
            'sh_signals': sh_signals
        }

    def validate_signals(self, signals: List[Dict], get_stock_data_func: Callable,
                         target_date: str, mark_as_buy: bool = False) -> List[Dict]:
        """
        验证信号是否满足跌幅条件，并可标记为买入信号
        """
        valid_signals = []
        for signal in tqdm(signals, desc="验证信号跌幅条件"):
            stock_code = signal.get('stock_code')
            market = signal.get('market', 'unknown')

            if self.check_decline_condition(stock_code, target_date, get_stock_data_func):
                signal['decline_condition_met'] = True
                valid_signals.append(signal)

                # 如果要求标记为买入信号
                if mark_as_buy:
                    self.mark_as_buy_signal(stock_code, target_date, market, signal)

        return valid_signals

    def check_decline_condition(self, stock_code: str, target_date: str,
                                get_stock_data_func: Callable) -> bool:
        """
        检查股票是否满足跌幅条件

        :param stock_code: 股票代码
        :param target_date: 目标日期
        :param get_stock_data_func: 获取股票数据的函数
        :return: 是否满足跌幅条件
        """
        try:
            # 获取最近3天数据 (target_date及前两个交易日)
            previous_days = self.trade_calendar.get_previous_trade_days(target_date, 2)
            if not previous_days:
                return False

            check_start = previous_days[0]
            stock_data = get_stock_data_func(stock_code, check_start, target_date)

            if len(stock_data) < 3:
                return False

            # 获取最近三天的收盘价
            closes = stock_data['close'].tail(3).values

            # 条件1: 连续三天下跌
            consecutive_decline = all(closes[i] < closes[i - 1] for i in range(1, 3))

            # 条件2: 三日累计跌幅超3%
            total_decline = (closes[-1] - closes[0]) / closes[0]
            decline_over_3pct = total_decline <= -0.03

            return consecutive_decline or decline_over_3pct
        except Exception as e:
            logging.error(f"检查跌幅条件失败 [{stock_code}]: {str(e)}")
            return False

    def mark_as_buy_signal(self, stock_code: str, signal_date: str, market: str, signal_data: Dict):
        """
        将信号标记为买入信号

        :param stock_code: 股票代码
        :param signal_date: 信号日期
        :param market: 市场类型 (sz/sh)
        :param signal_data: 信号数据
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        try:
            # 更新主信号表中的买入标记
            cursor.execute('''
            UPDATE stock_signals 
            SET buy_signal = 1, 
                decline_condition_met = 1
            WHERE stock_code = ? 
                AND signal_date = ? 
                AND market = ?
            ''', (stock_code, signal_date, market))

            # 添加到买入信号历史表
            cursor.execute('''
            INSERT INTO buy_signals_history 
            (stock_code, buy_date, signal_data)
            VALUES (?, ?, ?)
            ''', (stock_code, signal_date, json.dumps(signal_data)))

            conn.commit()
            logging.info(f"标记为买入信号: {stock_code} ({signal_date})")
        except Exception as e:
            logging.error(f"标记买入信号失败 [{stock_code}]: {str(e)}")
        finally:
            conn.close()

    def mark_as_sell_signal(self, stock_code: str, sell_date: str, buy_date: str, signal_data: Dict):
        """
        将信号标记为卖出信号

        :param stock_code: 股票代码
        :param sell_date: 卖出日期
        :param buy_date: 原始买入日期
        :param signal_data: 信号数据
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        try:
            # 更新主信号表中的卖出标记
            cursor.execute('''
            UPDATE stock_signals 
            SET sell_signal = 1
            WHERE stock_code = ? 
                AND signal_date = ? 
            ''', (stock_code, sell_date))

            # 添加到卖出信号历史表
            cursor.execute('''
            INSERT INTO sell_signals_history 
            (stock_code, sell_date, buy_date, signal_data)
            VALUES (?, ?, ?, ?)
            ''', (stock_code, sell_date, buy_date, json.dumps(signal_data)))

            conn.commit()
            logging.info(f"标记为卖出信号: {stock_code} (买入:{buy_date}, 卖出:{sell_date})")
        except Exception as e:
            logging.error(f"标记卖出信号失败 [{stock_code}]: {str(e)}")
        finally:
            conn.close()

    def get_active_buy_signals(self) -> List[Dict]:
        """
        获取所有尚未卖出的买入信号

        :return: 活跃买入信号列表
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        active_signals = []

        try:
            # 查询尚未卖出的买入信号
            cursor.execute('''
            SELECT b.stock_code, b.buy_date, b.signal_data
            FROM buy_signals_history b
            LEFT JOIN sell_signals_history s 
                ON b.stock_code = s.stock_code 
                AND b.buy_date = s.buy_date
            WHERE s.id IS NULL
            ''')

            for row in cursor.fetchall():
                signal = json.loads(row[2])
                signal['buy_date'] = row[1]
                active_signals.append(signal)

            logging.info(f"获取到 {len(active_signals)} 个活跃买入信号")
        except Exception as e:
            logging.error(f"获取活跃买入信号失败: {str(e)}")
        finally:
            conn.close()

        return active_signals

    def update_decline_conditions(self, target_date: str, get_stock_data_func: Callable):
        """
        更新指定日期所有信号的跌幅条件标记

        :param target_date: 目标日期
        :param get_stock_data_func: 获取股票数据的函数
        """
        # 获取目标日期的所有信号
        signals = self.get_signals_by_date(target_date)
        all_signals = signals['sz_signals'] + signals['sh_signals']

        updated_count = 0
        for signal in tqdm(all_signals, desc="更新跌幅条件标记"):
            stock_code = signal.get('stock_code')
            market = signal.get('market', 'unknown')

            condition_met = self.check_decline_condition(stock_code, target_date, get_stock_data_func)

            # 更新数据库标记
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            try:
                cursor.execute('''
                UPDATE stock_signals 
                SET decline_condition_met = ?
                WHERE stock_code = ? 
                    AND signal_date = ? 
                    AND market = ?
                ''', (1 if condition_met else 0, stock_code, target_date, market))
                conn.commit()
                updated_count += 1
            except Exception as e:
                logging.error(f"更新跌幅条件标记失败 [{stock_code}]: {str(e)}")
            finally:
                conn.close()

        logging.info(f"已更新 {updated_count} 条信号的跌幅条件标记")

    def fetch_and_generate_signals(self, target_date: str, batch_size: int = 100, days_offset: int = 180) -> Dict[
        str, Any]:
        """
        获取所有股票信号（多线程优化版）

        :param target_date: 目标日期 (YYYYMMDD)
        :param batch_size: 批量大小
        :param days_offset: 数据偏移天数
        :return: 包含信号的字典
        """
        # 确保stock_info_map已初始化
        if not self.stock_info_map:
            self._load_stock_info()

        # 获取日期范围
        try:
            start_date, end_date = self.trade_calendar.get_start_end_date(target_date, days_offset)
        except Exception as e:
            logging.warning(f"获取交易日期失败，使用默认起始日期: {str(e)}")
            start_date = "20240101"
            end_date = target_date

        # 处理深圳市场
        sz_signals = self._process_market_signals('sz', start_date, end_date, target_date)
        # 处理上海市场
        sh_signals = self._process_market_signals('sh', start_date, end_date, target_date)

        # 合并信号
        all_signals = sz_signals + sh_signals
        logging.info(f"信号生成完成: 总信号数 {len(all_signals)}")

        # 存储信号
        self.store_signals(sz_signals, 'sz')
        self.store_signals(sh_signals, 'sh')

        return {
            'sz_signals': sz_signals,
            'sh_signals': sh_signals,
            'target_date': target_date,
            'batch_size': batch_size,
            'days_offset': days_offset
        }

    def _load_stock_info(self):
        """加载股票基本信息"""
        try:

            stock_codes_df = self.stock_manager.get_stock_codes(market_types=['sz', 'sh'])
            self.stock_info_map = stock_codes_df.set_index('stock_code')[['stock_name', 'industry', 'rating']].to_dict(
                'index')
            logging.info(f"成功加载 {len(self.stock_info_map)} 条股票基本信息")
        except Exception as e:
            logging.error(f"加载股票基本信息失败: {str(e)}")
            self.stock_info_map = {}

    def _process_market_signals(self, market: str, start_date: str, end_date: str, target_date: str) -> List[Dict]:
        """
        处理指定市场的信号生成

        :param market: 市场类型 (sz/sh)
        :param start_date: 开始日期
        :param end_date: 结束日期
        :param target_date: 目标日期
        :return: 信号列表
        """
        try:
            # 获取市场股票代码
            stock_codes = self.stock_manager .get_stock_codes(market_types=[market])
            stock_codes_list = stock_codes['stock_code'].tolist()

            # 多线程获取股票数据
            stock_data_list = []
            with ThreadPoolExecutor(max_workers=10) as executor:
                future_to_code = {
                    executor.submit(self.data_manager.get_stock_data, code, start_date, end_date): code
                    for code in stock_codes_list
                }

                with tqdm(total=len(future_to_code), desc=f"获取{market}股票数据") as pbar:
                    for future in as_completed(future_to_code):
                        try:
                            stock_data = future.result()
                            if not stock_data.empty and 'date' in stock_data.columns:
                                stock_data = stock_data.sort_values('date')
                                stock_code = future_to_code[future]

                                # 添加股票基本信息
                                stock_data['stock_name'] = self.stock_info_map.get(stock_code, {}).get('stock_name', '未知')
                                stock_data['industry'] = self.stock_info_map.get(stock_code, {}).get('industry', '未知')
                                stock_data['rating'] = self.stock_info_map.get(stock_code, {}).get('rating', '未知')

                                stock_data_list.append(stock_data)
                        except Exception as e:
                            logging.error(f"处理股票数据失败: {str(e)}")
                        pbar.update(1)

            # 批量生成信号
            return self.batch_generate_signals(stock_data_list, target_date)
        except Exception as e:
            logging.error(f"处理{market}市场信号失败: {str(e)}")
            return []

    def batch_generate_signals(self, stock_data_list: List[pd.DataFrame], target_date: str) -> List[Dict]:
        """
        批量生成信号

        :param stock_data_list: 股票数据列表
        :param target_date: 目标日期 (YYYYMMDD)
        :return: 信号列表
        """
        if not stock_data_list:
            return []

        # 转换为目标日期格式
        target_dt = datetime.strptime(target_date, '%Y%m%d')
        target_date_obj = target_dt.date()

        all_signals = []
        for i in range(0, len(stock_data_list), self.batch_size):
            batch = stock_data_list[i:i + self.batch_size]
            try:
                batch_result = self.signal_generator.generate_signals(batch)

                # 过滤目标日期的信号
                for signal in batch_result.get('signals_list', []):
                    try:
                        signal_date = pd.to_datetime(signal['date']).date()
                        if signal_date == target_date_obj:
                            # 统一转换为YYYYMMDD格式
                            signal['date'] = signal_date.strftime('%Y%m%d')
                            all_signals.append(signal)
                    except Exception as e:
                        logging.warning(f"信号日期转换失败: {str(e)}")
                        continue
            except Exception as e:
                logging.error(f"批量生成信号失败: {str(e)}")

        return all_signals


# 示例使用代码
if __name__ == "__main__":
    # 初始化
    trade_calendar = TradeDateManager()
    data_manager = YearlyStockDataManager()
    signal_generator = SignalGenerator()

    stock_manager = StockBasicInfoManager()
    signal_manager = StockSignalsManager(
        trade_calendar=trade_calendar,
        data_manager=data_manager,
        signal_generator=signal_generator,
        stock_manager=stock_manager
    )

    # 生成并存储信号
    target_date = "20230630"
    signals_result = signal_manager.fetch_and_generate_signals(target_date)

    # 验证信号
    valid_signals = signal_manager.validate_signals(
        signals_result['sz_signals'] + signals_result['sh_signals'],
        data_manager.get_stock_data,
        target_date,
        mark_as_buy=True
    )

    # 获取买入信号
    buy_signals = signal_manager.get_signals_by_date(target_date, filter_condition='buy')

    # 模拟卖出条件判断
    def should_sell(buy_signal):
        # 这里实现卖出条件判断逻辑
        return True  # 示例：总是卖出

    # 标记卖出信号
    current_date = "20230701"
    for buy_signal in signal_manager.get_active_buy_signals():
        if should_sell(buy_signal):
            signal_manager.mark_as_sell_signal(
                buy_signal['stock_code'],
                current_date,
                buy_signal['buy_date'],
                buy_signal
            )

    # 更新跌幅条件标记
    signal_manager.update_decline_conditions("20230701", data_manager.get_stock_data)
